# Copyright(C) 2023. Huawei Technologies Co.,Ltd. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import argparse
import os
from multiprocessing import Process, Queue

import cv2

from wav2lip import Wav2lip
from wav2lip_utils.audio_utils import load_wav
from wav2lip_utils.constant import *


def parse_arguments() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-o",
        "--output_dir",
        type=str,
        default="./results",
        help="path of directory to save results.",
    )
    parser.add_argument(
        "-c",
        "--config",
        type=str,
        default="./config.yaml",
        help="Path of config file.",
    )

    return parser.parse_args()


class temp():
    def set(self):
        pass


if __name__ == '__main__':
    args = parse_arguments()
    if not os.path.exists(args.output_dir):
        os.makedirs(args.output_dir, mode=0o750, exist_ok=True)

    in_q = Queue(QUEUE_SIZE)
    out_q = Queue(QUEUE_SIZE)
    wav2lip_cls = Wav2lip(args.config)
    wav2lip_proc = Process(target=wav2lip_cls.run, args=(in_q, out_q, temp()))
    wav2lip_proc.start()

    input_wav = load_wav('./assets/test_audio.wav', W2L_SAMPLE_RATE)
    in_q.put((input_wav, True))

    while True:
        if not out_q.empty():
            result = out_q.get()[0]
            if result is None:
                break
            else:
                print(len(result))
                for i, frame in enumerate(result):
                    cv2.imwrite(os.path.join(args.output_dir, f'{i}.jpg'), frame)
                print('---------------------')
